#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/2/28 19:58
@desc: 用于生成 tensorflow 支持的二进制文件
"""
import os
from collections import Iterable

import numpy as np
import tensorflow as tf

from dltools.utils import log


class BaseRecordGenerator(object):
    """
    用于生成 tensorflow 支持的二进制文件的基本类
    """

    def __init__(self, data, output,
                 mean_value_length=0,
                 display=100,
                 logger=None):
        """

        Parameters
        ----------
        data: 可迭代的数据对象
        output: 保存 record 的文件名
        mean_value_length: 是否计算数据均值, 为 0 不计算，大于 0 时数字就是均值长度
        display: 显示间隔
        logger: 日志对象
        """
        assert isinstance(data, Iterable), '请输入正确的数据！data必须是可迭代的对象。'
        self.data = data

        self.mean_value_length = mean_value_length
        if self.mean_value_length == 0:
            self.mean = np.zeros((self.mean_value_length, ), dtype=np.float128)
        else:
            self.mean = None

        if not os.path.exists(os.path.dirname(output)):
            os.makedirs(os.path.dirname(output))

        self.writer = tf.python_io.TFRecordWriter(output)

        self.buf_data = {}
        self.display = display
        self.mean = None
        self.count = 0

        if logger is not None:
            self.logger = logger
        else:
            self.logger = log.get_console_logger('RecordGenerator')

    def _encode_data(self):
        """
        将数据转化为 tf example

        Returns
        -------

        """
        raise NotImplementedError

    def _write_data(self):
        """
        将 tf example 写入 record 文件

        Returns
        -------

        """
        self.writer.write(self.buf_data['example'].SerializeToString())

    def update(self):
        """
        处理全部数据

        Returns
        -------

        """
        for meta in self.data:
            if self.count % self.display == 0:
                self.logger.info(
                    'Processing the number of {} data.'.format(self.count))
            self.buf_data['raw'] = meta
            self._encode_data()
            self._write_data()
            self.buf_data.clear()
            self.count += 1
        if self.mean_value_length:
            self.mean = self.mean / self.count
        self.writer.close()
